import tensorflow as tf
import numpy as np
import pandas as pd
import logging
from sklearn.model_selection import train_test_split
from sklearn.utils import shuffle
import joblib
from typing import Tuple, Dict, Optional

from ..utils.data_generator import create_tf_dataset

class HybridTrainer:
    """混合模型训练器 - 完整版本"""
    
    def __init__(self, model, config=None):
        self.model = model
        self.config = config or {}
        self.logger = logging.getLogger(__name__)
        self.history = None
        self.scalers = {}
    
    def prepare_training_data(self, ratings_df: pd.DataFrame, 
                            user_features: pd.DataFrame, 
                            movie_features: pd.DataFrame,
                            user_id_map: Dict, 
                            movie_id_map: Dict,
                            test_size: float = 0.2,
                            sample_ratio: float = 1.0) -> Tuple:
        """
        准备训练数据 - 支持大数据集抽样
        """
        self.logger.info("准备训练数据...")
        
        # 数据抽样（对于大数据集）
        if sample_ratio < 1.0:
            original_size = len(ratings_df)
            ratings_df = ratings_df.sample(frac=sample_ratio, random_state=42)
            self.logger.info(f"数据抽样: {original_size} -> {len(ratings_df)}")
        
        # 映射ID
        ratings_df = ratings_df.copy()
        ratings_df['user_idx'] = ratings_df['userId'].map(user_id_map)
        ratings_df['movie_idx'] = ratings_df['movieId'].map(movie_id_map)
        
        # 移除无法映射的记录
        before_size = len(ratings_df)
        ratings_df = ratings_df.dropna(subset=['user_idx', 'movie_idx'])
        after_size = len(ratings_df)
        if before_size != after_size:
            self.logger.warning(f"移除无法映射的记录: {before_size} -> {after_size}")

        # 再次确保 movie_idx 在合法范围内，防止 embedding 查表越界
        n_movies = len(movie_id_map)
        before_size2 = len(ratings_df)
        ratings_df = ratings_df[(ratings_df['movie_idx'] >= 0) & (ratings_df['movie_idx'] < n_movies)]
        after_size2 = len(ratings_df)
        if before_size2 != after_size2:
            self.logger.warning(f"移除 movie_idx 越界的记录: {before_size2} -> {after_size2}")
        
        # 准备特征数据
        user_indices = ratings_df['user_idx'].values.astype(np.int32)
        movie_indices = ratings_df['movie_idx'].values.astype(np.int32)
        
        # 获取用户特征（确保索引对齐，使用.loc更高效）
        user_feats = user_features.reindex(ratings_df['userId']).fillna(0).values.astype(np.float32)
        # 确保 movie_features 的 index 是 movieId，且无重复，只保留数值型特征
        if 'movieId' in movie_features.columns:
            movie_features = movie_features.drop_duplicates(subset='movieId').set_index('movieId')
        numeric_cols = movie_features.select_dtypes(include=[np.number]).columns
        movie_features = movie_features[numeric_cols]

        ratings = ratings_df['rating'].values.astype(np.float32)
        # movie_ids = ratings_df['movieId'].values.astype(np.int32)  # 不再作为 embedding 输入

        # 数据统计
        self.logger.info(f"训练数据统计:")
        self.logger.info(f"  用户索引: {user_indices.shape}")
        self.logger.info(f"  电影索引: {movie_indices.shape}")
        self.logger.info(f"  用户特征: {user_feats.shape}")
        self.logger.info(f"  评分: {ratings.shape}")

        # 分割训练集和测试集
        if test_size > 0:
            (user_indices_train, user_indices_test,
             movie_indices_train, movie_indices_test,
             user_feats_train, user_feats_test,
             ratings_train, ratings_test) = train_test_split(
                user_indices, movie_indices, user_feats, ratings,
                test_size=test_size, random_state=42, stratify=None
            )

            train_data = (user_indices_train, movie_indices_train, user_feats_train, ratings_train)
            test_data = (user_indices_test, movie_indices_test, user_feats_test, ratings_test)
            self.movie_features = movie_features
            return train_data, test_data, movie_features
        else:
            self.movie_features = movie_features
            return (user_indices, movie_indices, user_feats, ratings), None, movie_features
    
    def train(self, train_data, val_data=None, epochs=50, batch_size=1024, 
              use_generator=False, callbacks=None):
        """训练模型"""
        self.logger.info("开始模型训练...")

        # 准备输入数据
        (user_idx_train, movie_idx_train, user_feat_train, y_train) = train_data
        
        # 将输入数据转换为字典格式
        X_train = {
            "user_id_input": user_idx_train,
            "movie_id_input": movie_idx_train,
            "user_feature_input": user_feat_train
        }
        
        # 使用 tf.data.Dataset 创建高效的数据管道
        self.logger.info("使用 tf.data.Dataset 创建高效的数据管道...")
        train_dataset = create_tf_dataset(X_train, y_train, batch_size, shuffle=True)

        val_dataset = None
        if val_data is not None:
            (user_idx_val, movie_idx_val, user_feat_val, y_val) = val_data
            X_val = {
                "user_id_input": user_idx_val,
                "movie_id_input": movie_idx_val,
                "user_feature_input": user_feat_val
            }
            val_dataset = create_tf_dataset(X_val, y_val, batch_size, shuffle=False)

        # 训练配置
        if callbacks is None:
            callbacks = self.model.get_callbacks() if hasattr(self.model, 'get_callbacks') else []
        
        # 直接使用 tf.data.Dataset 进行训练
        self.history = self.model.fit(
            train_dataset,
            epochs=epochs,
            validation_data=val_dataset,
            callbacks=callbacks,
            verbose=1
        )
        
        self.logger.info("模型训练完成")
        return self.history
    
    def evaluate(self, test_data):
        """评估模型"""
        if test_data is None:
            self.logger.warning("没有测试数据可供评估")
            return None
        
        (user_idx_test, movie_idx_test, user_feat_test, y_test) = test_data
        
        X_test = {
            "user_id_input": user_idx_test,
            "movie_id_input": movie_idx_test,
            "user_feature_input": user_feat_test
        }
        
        evaluation = self.model.evaluate(X_test, y_test, verbose=0)
        
        metrics = {}
        # The model's metrics_names might not be populated until after training/evaluation
        # so we'll manually name them based on the model's compile configuration.
        metric_names = ['loss'] + [m.name for m in self.model.metrics]
        
        # Ensure evaluation is a list
        if not isinstance(evaluation, list):
            evaluation = [evaluation]
            
        for i, metric_name in enumerate(metric_names):
            if i < len(evaluation):
                 metrics[metric_name] = evaluation[i]

        self.logger.info("模型评估结果:")
        for metric, value in metrics.items():
            self.logger.info(f"  {metric}: {value:.4f}")
        
        return metrics
    
    def predict_rating(self, user_id, movie_id, user_features, movie_features, 
                      user_id_map, movie_id_map):
        """预测单个评分"""
        try:
            user_idx = user_id_map.get(user_id, -1)
            movie_idx = movie_id_map.get(movie_id, -1)
            
            if user_idx == -1 or movie_idx == -1:
                self.logger.warning(f"无法找到用户 {user_id} 或电影 {movie_id} 的映射")
                return None
            
            # 准备输入数据
            user_idx = np.array([user_idx])
            movie_idx = np.array([movie_idx])
            user_feat = np.array([user_features])
            movie_feat = np.array([movie_features])
            
            prediction = self.model.predict(
                [user_idx, movie_idx, user_feat, movie_feat],
                verbose=0
            )
            
            return float(prediction[0][0])
        
        except Exception as e:
            self.logger.error(f"预测错误: {e}")
            return None
    
    def save_training_info(self, filepath):
        """保存训练信息"""
        training_info = {
            'history': self.history.history if self.history else {},
            'config': self.config,
            'scalers': self.scalers
        }
        joblib.dump(training_info, filepath)
        self.logger.info(f"训练信息已保存到: {filepath}")